# -*- coding: utf-8 -*-
import numpy as np
import os, cv2, time, copy, random
import imgaug.augmenters as iaa

def increase_saturation(image, factor=50):
    """
    增加图像的饱和度

    参数:
        image: 输入图像 (numpy array)，应为BGR格式
        factor: 饱和度增加的因子，大于1表示增加饱和度

    返回:
        result_image: 饱和度增加后的图像
    """
    # 将图像从BGR色彩空间转换到HSV色彩空间
    hsv_image = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)

    # 分离H,S,V三个通道
    h, s, v = cv2.split(hsv_image)

    # 增加饱和度，注意避免溢出[0, 255]
    lim = 255 - factor
    s[s > lim] = 255
    s[s <= lim] += factor

    # 合并处理后的H,S,V通道
    final_hsv = cv2.merge((h, s, v))

    # 将图像从HSV色彩空间转换回BGR色彩空间
    result_image = cv2.cvtColor(final_hsv, cv2.COLOR_HSV2BGR)

    return result_image


def enhance_color_clahe(image):
    # 将图像从 BGR 转换到 LAB 色彩空间
    lab = cv2.cvtColor(image, cv2.COLOR_BGR2LAB)
    l, a, b = cv2.split(lab)

    # 应用 CLAHE 到 L 通道
    clahe = cv2.createCLAHE(clipLimit=3.0, tileGridSize=(8, 8))
    cl = clahe.apply(l)

    # 合并通道
    limg = cv2.merge((cl,a,b))

    # 转换回 BGR 色彩空间
    return cv2.cvtColor(limg, cv2.COLOR_LAB2BGR)

def image2pdf(image, gamma=1.2):
    image = increase_saturation(image, factor=10)
    # image = enhance_color_clahe(image)

    # # 应用查找表进行Gamma校正
    # invGamma = 1.0 / gamma
    # table = np.array([((i / 255.0) ** invGamma) * 255 for i in np.arange(0, 256)]).astype("uint8")
    # gamma_corrected = cv2.LUT(image, table)

    gamma = (1.2, 1.3)
    seq = iaa.GammaContrast(gamma)
    aug_light_image = seq(image=image)

    # 双边滤波
    result = cv2.bilateralFilter(aug_light_image, d=23, sigmaColor=15, sigmaSpace=35)  # 尽可能保留轮廓清晰
    return result


def pdf2image_by_bright(image,
                        dpi=2000,
                        screen_angle=45,
                        noise_strength=0.1,
                        color_noise_level=30,
                        brightness_adjustment=40):

    """
    在彩色图像上添加轻微的半色调网点噪声，并在此基础上添加细微的颜色噪声和亮度调整。

    :param image: 输入彩色图像 (BGR 格式)
    :param dpi: 输出图像的分辨率，影响网点大小
    :param screen_angle: 半色调网点的角度
    :param noise_strength: 网点噪声的强度 (0-1)，值越小噪声越轻微
    :param color_noise_level: 颜色噪声的水平 (0-255)，值越大噪声越明显
    :param brightness_adjustment: 亮度调整量 (-255 到 255)，正值使图像更亮，负值使图像更暗
    :return: 添加了轻微网点噪声、颜色噪声及亮度调整后的彩色图像
    """
    # 将图像从 BGR 转换为 RGB 格式
    rgb_image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

    # 分离 RGB 通道
    r_channel, g_channel, b_channel = cv2.split(rgb_image)

    def halftone_channel(channel, angle):
        height, width = channel.shape

        # 计算网点大小和角度变换矩阵
        cell_size = max(1, int(72 / dpi * 25.4))  # 根据 DPI 计算单元格大小

        angle_rad = np.deg2rad(angle)
        x_freq = np.cos(angle_rad) / cell_size
        y_freq = np.sin(angle_rad) / cell_size

        # 创建网格
        x_grid, y_grid = np.meshgrid(np.arange(width), np.arange(height))

        # 计算相位
        phase = np.sin(2 * np.pi * (x_freq * x_grid + y_freq * y_grid))

        # 将相位映射到 [0, 255] 范围
        halftone_noise = ((phase + 1) / 2 * 255).astype(np.uint8)

        return halftone_noise

    # 应用不同的屏幕角度到每个通道
    halftone_r = halftone_channel(r_channel, screen_angle)
    halftone_g = halftone_channel(g_channel, screen_angle + 30)
    halftone_b = halftone_channel(b_channel, screen_angle + 60)

    # 合并网点噪声通道
    halftone_noise_image = cv2.merge([halftone_b, halftone_g, halftone_r])

    # 将网点噪声与原始图像混合
    halftone_noise_image = cv2.cvtColor(halftone_noise_image, cv2.COLOR_RGB2BGR)
    noisy_image = cv2.addWeighted(image, 1 - noise_strength, halftone_noise_image, noise_strength, 0)

    # 为每个像素添加轻微的颜色噪声
    h, w, _ = noisy_image.shape
    color_noise = np.random.randint(-color_noise_level, color_noise_level, (h, w, 3), dtype='int16')
    noisy_image = noisy_image.astype('int16') + color_noise
    noisy_image = np.clip(noisy_image, 0, 255).astype('uint8')

    # 调整亮度
    if brightness_adjustment != 0:
        noisy_image = noisy_image.astype('int16')  # 转换为 int16 以避免溢出
        noisy_image += brightness_adjustment
        noisy_image = np.clip(noisy_image, 0, 255).astype('uint8')  # 确保像素值在 [0, 255] 范围内

    result = cv2.GaussianBlur(noisy_image, (5, 5), 0)
    return result


def pdf2image_by_contrast(image,
                          dpi=2000,
                          screen_angle=45,
                          noise_strength=0.1,
                          color_noise_level=30,
                          hue_variation_range=(-10, 10),
                          contrast_variation_range=(0.8, 1.2)):
    """
    在彩色图像上添加轻微的半色调网点噪声，并在此基础上添加细微的颜色噪声、色调变化和对比度变化。

    :param image: 输入彩色图像 (BGR 格式)
    :param dpi: 输出图像的分辨率，影响网点大小
    :param screen_angle: 半色调网点的角度
    :param noise_strength: 网点噪声的强度 (0-1)，值越小噪声越轻微
    :param color_noise_level: 颜色噪声的水平 (0-255)，值越大噪声越明显
    :param hue_variation_range: 色调变化范围 (min, max)，单位为度 (-180 到 180)
    :param contrast_variation_range: 对比度变化范围 (min_alpha, max_alpha)
    :return: 添加了轻微网点噪声、颜色噪声、色调变化和对比度变化后的彩色图像
    """
    # 将图像从 BGR 转换为 RGB 格式
    rgb_image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

    # 分离 RGB 通道
    r_channel, g_channel, b_channel = cv2.split(rgb_image)

    def halftone_channel(channel, angle):
        height, width = channel.shape

        # 计算网点大小和角度变换矩阵
        cell_size = max(1, int(72 / dpi * 25.4))  # 根据 DPI 计算单元格大小

        angle_rad = np.deg2rad(angle)
        x_freq = np.cos(angle_rad) / cell_size
        y_freq = np.sin(angle_rad) / cell_size

        # 创建网格
        x_grid, y_grid = np.meshgrid(np.arange(width), np.arange(height))

        # 计算相位
        phase = np.sin(2 * np.pi * (x_freq * x_grid + y_freq * y_grid))

        # 将相位映射到 [0, 255] 范围
        halftone_noise = ((phase + 1) / 2 * 255).astype(np.uint8)

        return halftone_noise

    # 应用不同的屏幕角度到每个通道
    halftone_r = halftone_channel(r_channel, screen_angle)
    halftone_g = halftone_channel(g_channel, screen_angle + 30)
    halftone_b = halftone_channel(b_channel, screen_angle + 60)

    # 合并网点噪声通道
    halftone_noise_image = cv2.merge([halftone_b, halftone_g, halftone_r])

    # 将网点噪声与原始图像混合
    halftone_noise_image = cv2.cvtColor(halftone_noise_image, cv2.COLOR_RGB2BGR)
    noisy_image = cv2.addWeighted(image, 1 - noise_strength, halftone_noise_image, noise_strength, 0)

    # 为每个像素添加轻微的颜色噪声
    h, w, _ = noisy_image.shape
    color_noise = np.random.randint(-color_noise_level, color_noise_level, (h, w, 3), dtype='int16')
    noisy_image = noisy_image.astype('int16') + color_noise
    noisy_image = np.clip(noisy_image, 0, 255).astype('uint8')

    # 转换到 HSV 颜色空间以调整色调
    hsv_image = cv2.cvtColor(noisy_image, cv2.COLOR_BGR2HSV)

    # 分离 HSV 通道
    h_channel, s_channel, v_channel = cv2.split(hsv_image)

    # 在指定范围内随机调整色调
    hue_min, hue_max = hue_variation_range
    h_channel = (h_channel.astype(np.int16) + np.random.randint(hue_min, hue_max + 1, h_channel.shape)) % 180
    h_channel = h_channel.astype(np.uint8)

    # 合并调整后的 HSV 通道
    adjusted_hsv_image = cv2.merge([h_channel, s_channel, v_channel])

    # 转换回 BGR 颜色空间
    adjusted_image = cv2.cvtColor(adjusted_hsv_image, cv2.COLOR_HSV2BGR)

    # 随机生成对比度因子
    contrast_min, contrast_max = contrast_variation_range
    alpha = np.random.uniform(contrast_min, contrast_max)  # 随机生成对比度因子
    beta = 0  # 不调整亮度

    # 应用对比度变换
    adjusted_image = cv2.convertScaleAbs(adjusted_image, alpha=alpha, beta=beta)
    result = cv2.GaussianBlur(adjusted_image, (5, 5), 0)
    return result


def pdf2image_by_hue(image,
              dpi=2000,
              screen_angle=45,
              noise_strength=0.1,
              color_noise_level=30,
              hue_variation_range=(-10, 10)):
    """
    在彩色图像上添加轻微的半色调网点噪声，并在此基础上添加细微的颜色噪声和色调变化。

    :param image: 输入彩色图像 (BGR 格式)
    :param dpi: 输出图像的分辨率，影响网点大小
    :param screen_angle: 半色调网点的角度
    :param noise_strength: 网点噪声的强度 (0-1)，值越小噪声越轻微
    :param color_noise_level: 颜色噪声的水平 (0-255)，值越大噪声越明显
    :param hue_variation_range: 色调变化范围 (min, max)，单位为度 (-180 到 180)
    :return: 添加了轻微网点噪声、颜色噪声及色调变化后的彩色图像
    """
    # 将图像从 BGR 转换为 RGB 格式
    rgb_image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

    # 分离 RGB 通道
    r_channel, g_channel, b_channel = cv2.split(rgb_image)

    def halftone_channel(channel, angle):
        height, width = channel.shape

        # 计算网点大小和角度变换矩阵
        cell_size = max(1, int(72 / dpi * 25.4))  # 根据 DPI 计算单元格大小

        angle_rad = np.deg2rad(angle)
        x_freq = np.cos(angle_rad) / cell_size
        y_freq = np.sin(angle_rad) / cell_size

        # 创建网格
        x_grid, y_grid = np.meshgrid(np.arange(width), np.arange(height))

        # 计算相位
        phase = np.sin(2 * np.pi * (x_freq * x_grid + y_freq * y_grid))

        # 将相位映射到 [0, 255] 范围
        halftone_noise = ((phase + 1) / 2 * 255).astype(np.uint8)

        return halftone_noise

    # 应用不同的屏幕角度到每个通道
    halftone_r = halftone_channel(r_channel, screen_angle)
    halftone_g = halftone_channel(g_channel, screen_angle + 30)
    halftone_b = halftone_channel(b_channel, screen_angle + 60)

    # 合并网点噪声通道
    halftone_noise_image = cv2.merge([halftone_b, halftone_g, halftone_r])

    # 将网点噪声与原始图像混合
    halftone_noise_image = cv2.cvtColor(halftone_noise_image, cv2.COLOR_RGB2BGR)
    noisy_image = cv2.addWeighted(image, 1 - noise_strength, halftone_noise_image, noise_strength, 0)

    # 为每个像素添加轻微的颜色噪声
    h, w, _ = noisy_image.shape
    color_noise = np.random.randint(-color_noise_level, color_noise_level, (h, w, 3), dtype='int16')
    noisy_image = noisy_image.astype('int16') + color_noise
    noisy_image = np.clip(noisy_image, 0, 255).astype('uint8')

    # 转换到 HSV 颜色空间以调整色调
    hsv_image = cv2.cvtColor(noisy_image, cv2.COLOR_BGR2HSV)

    # 分离 HSV 通道
    h_channel, s_channel, v_channel = cv2.split(hsv_image)

    # 在指定范围内随机调整色调
    hue_min, hue_max = hue_variation_range
    h_channel = (h_channel.astype(np.int16) + np.random.randint(hue_min, hue_max + 1, h_channel.shape)) % 180
    h_channel = h_channel.astype(np.uint8)

    # 合并调整后的 HSV 通道
    adjusted_hsv_image = cv2.merge([h_channel, s_channel, v_channel])

    # 转换回 BGR 颜色空间
    adjusted_image = cv2.cvtColor(adjusted_hsv_image, cv2.COLOR_HSV2BGR)

    result = cv2.GaussianBlur(adjusted_image, (5, 5), 0)
    return result


